import librosa
import numpy
from scipy.io import wavfile as wav
from python_speech_features import fbank
import noisereduce as nr
import numpy as np

from .config import *


def volume_normalization(signal):
    # 检查原始信号的数据类型
    original_dtype = signal.dtype

    # 将信号转换为浮点类型进行标准化
    float_signal = signal.astype(np.float32)
    normalized_signal = librosa.util.normalize(float_signal)

    # 根据原始数据类型转换回相应的整数类型
    if original_dtype == np.int16:
        max_int16 = np.iinfo(np.int16).max
        return (normalized_signal * max_int16).astype(np.int16)
    elif original_dtype == np.int32:
        max_int32 = np.iinfo(np.int32).max
        return (normalized_signal * max_int32).astype(np.int32)
    elif original_dtype == np.int8:
        max_int8 = np.iinfo(np.int8).max
        return (normalized_signal * max_int8).astype(np.int8)
    else:
        raise ValueError(f"Unsupported data type: {original_dtype}")

def split_audio(signal, rate, split_duration=1.3, top_db=60):
    """
    使用静音检测分割音频，并返回长度大于min_duration和split_duration的非静音段。
    signal: 音频信号
    rate: 采样率
    split_duration: 最短有效片段长度（秒）
    top_db: 分割的静音阈值（dB）
    """
    # 检测非静音部分
    non_silent_intervals = librosa.effects.split(signal, top_db=top_db)
    valid_segments = []
    for start, end in non_silent_intervals:
        # 将非静音部分按 split_duration 划分
        segment_start = int(start)
        while segment_start < int(end):
            segment_end = int(min(segment_start + int(split_duration * rate), end))
            if (segment_end - segment_start) / rate >= split_duration:
                segment = signal[segment_start:segment_end]
                if np.issubdtype(segment.dtype, np.floating):
                    segment = (segment * 32767).astype(np.int16)
                valid_segments.append(segment)
            segment_start = segment_end
    return valid_segments

def extract_features(signal, rate):
    # 分割音频并提取特征
    audio_segments = split_audio(signal, rate)
    print(signal, rate)
    fbank_feats = []
    for segment in audio_segments:
        # 噪声消除
        reduced_segment = nr.reduce_noise(y=segment, sr=rate)
        # 音量标准化
        normalized_segment = volume_normalization(reduced_segment)

        print(segment)
        print(normalized_segment)
        # 提取 fbank 特征
        features, energy = fbank(normalized_segment, samplerate=rate, nfilt=26, winfunc=np.hamming)
        # 对特征进行对数变换
        log_features = np.log(features)
        fbank_feats.append(log_features)
    return fbank_feats
